from sqlalchemy import create_engine, MetaData, Table, inspect
from sqlalchemy.exc import SQLAlchemyError
from typing import List, Dict, Any
import logging
from models import DatabaseConnection, TableInfo, TableColumn, Relationship

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class DatabaseAnalyzer:
    def __init__(self, connection: DatabaseConnection):
        self.connection = connection
        self.engine = None
        self.metadata = None
        self._connect()
    
    def _connect(self):
        """建立数据库连接"""
        try:
            if self.connection.db_type == "mysql":
                connection_string = f"mysql+pymysql://{self.connection.username}:{self.connection.password}@{self.connection.host}:{self.connection.port}/{self.connection.database}?charset={self.connection.charset}"
            elif self.connection.db_type == "postgresql":
                connection_string = f"postgresql+psycopg2://{self.connection.username}:{self.connection.password}@{self.connection.host}:{self.connection.port}/{self.connection.database}"
            elif self.connection.db_type == "sqlserver":
                connection_string = f"mssql+pyodbc://{self.connection.username}:{self.connection.password}@{self.connection.host}:{self.connection.port}/{self.connection.database}?driver=ODBC+Driver+17+for+SQL+Server"
            elif self.connection.db_type == "sqlite":
                connection_string = f"sqlite:///{self.connection.database}"
            else:
                raise ValueError(f"不支持的数据库类型: {self.connection.db_type}")
            
            self.engine = create_engine(connection_string)
            self.metadata = MetaData()
            logger.info(f"成功连接到数据库: {self.connection.database}")
            
        except Exception as e:
            logger.error(f"数据库连接失败: {str(e)}")
            raise
    
    def get_tables(self) -> List[TableInfo]:
        """获取所有表结构信息"""
        try:
            inspector = inspect(self.engine)
            table_names = inspector.get_table_names()
            
            tables = []
            for table_name in table_names:
                table_info = self._get_table_info(inspector, table_name)
                tables.append(table_info)
            
            return tables
            
        except Exception as e:
            logger.error(f"获取表结构失败: {str(e)}")
            raise
    
    def _get_table_info(self, inspector, table_name: str) -> TableInfo:
        """获取单个表的详细信息"""
        try:
            # 获取列信息
            columns = inspector.get_columns(table_name)
            
            # 获取主键
            primary_keys = inspector.get_pk_constraint(table_name)
            pk_columns = primary_keys.get('constrained_columns', []) if primary_keys else []
            
            # 获取外键
            foreign_keys = inspector.get_foreign_keys(table_name)
            
            # 获取表注释（SQLite可能不支持）
            comment = None
            try:
                if self.connection.db_type != "sqlite":
                    table_comment = inspector.get_table_comment(table_name)
                    comment = table_comment.get('text', '') if table_comment else None
            except:
                comment = None
            
            # 构建列信息
            table_columns = []
            for column in columns:
                try:
                    # 获取列注释（如果支持）
                    column_comment = None
                    try:
                        if self.connection.db_type != "sqlite" and 'comment' in column:
                            column_comment = column.get('comment')
                    except:
                        column_comment = None
                    
                    column_info = TableColumn(
                        name=column['name'],
                        data_type=str(column['type']),
                        is_nullable=column.get('nullable', True),
                        is_primary_key=column['name'] in pk_columns,
                        is_foreign_key=self._is_foreign_key(column['name'], foreign_keys),
                        default_value=str(column['default']) if column.get('default') else None,
                        comment=column_comment
                    )
                    
                    # 设置外键引用信息
                    if column_info.is_foreign_key:
                        for fk in foreign_keys:
                            if column['name'] in fk['constrained_columns']:
                                column_info.referenced_table = fk['referred_table']
                                column_info.referenced_column = fk['referred_columns'][0] if fk['referred_columns'] else None
                                break
                    
                    table_columns.append(column_info)
                    
                except Exception as col_error:
                    logger.warning(f"处理列 {column['name']} 时出错: {str(col_error)}")
                    continue
            
            # 构建外键信息
            fk_info = []
            for fk in foreign_keys:
                fk_info.append({
                    'constrained_columns': fk['constrained_columns'],
                    'referred_table': fk['referred_table'],
                    'referred_columns': fk['referred_columns']
                })
            
            return TableInfo(
                name=table_name,
                comment=comment,
                columns=table_columns,
                primary_keys=pk_columns,
                foreign_keys=fk_info
            )
            
        except Exception as e:
            logger.error(f"获取表 {table_name} 信息失败: {str(e)}")
            raise
    
    def _is_foreign_key(self, column_name: str, foreign_keys: List[Dict]) -> bool:
        """判断列是否为外键"""
        for fk in foreign_keys:
            if column_name in fk['constrained_columns']:
                return True
        return False
    
    def _is_column_unique(self, inspector, table_name: str, column_name: str) -> bool:
        """检查列是否有唯一约束"""
        try:
            # 获取唯一约束
            unique_constraints = inspector.get_unique_constraints(table_name)
            for constraint in unique_constraints:
                if column_name in constraint['column_names']:
                    return True
            
            # 检查是否为主键（主键默认唯一）
            pk_constraint = inspector.get_pk_constraint(table_name)
            if pk_constraint and column_name in pk_constraint.get('constrained_columns', []):
                return True
            
            return False
        except:
            return False
    
    def _is_many_to_many_table(self, inspector, table_name: str) -> bool:
        """判断表是否可能是多对多关系表"""
        try:
            # 多对多关系表的启发式判断：
            # 1. 表名包含关联词汇
            # 2. 有多个外键
            # 3. 列数较少（通常只包含外键和少量附加字段）
            
            foreign_keys = inspector.get_foreign_keys(table_name)
            columns = inspector.get_columns(table_name)
            
            # 检查表名是否包含关联词汇
            association_words = ['link', 'relation', 'association', 'bridge', 'junction', 'map', 'mapping']
            is_association_name = any(word in table_name.lower() for word in association_words)
            
            # 有多个外键且列数较少
            has_multiple_fks = len(foreign_keys) >= 2
            has_few_columns = len(columns) <= 5  # 通常多对多表列数较少
            
            return is_association_name or (has_multiple_fks and has_few_columns)
        except:
            return False
    
    def get_relationships(self) -> List[Relationship]:
        """获取表之间的关系"""
        try:
            inspector = inspect(self.engine)
            table_names = inspector.get_table_names()
            
            relationships = []
            for table_name in table_names:
                foreign_keys = inspector.get_foreign_keys(table_name)
                
                for fk in foreign_keys:
                    referred_table = fk['referred_table']
                    constrained_columns = fk['constrained_columns']
                    referred_columns = fk['referred_columns']
                    
                    for i, constrained_column in enumerate(constrained_columns):
                        referred_column = referred_columns[i] if i < len(referred_columns) else referred_columns[0]
                        
                        # 检查是否为唯一约束（可能是一对一关系）
                        is_unique = self._is_column_unique(inspector, table_name, constrained_column)
                        
                        # 检查是否可能是多对多关系
                        is_many_to_many = self._is_many_to_many_table(inspector, table_name)
                        
                        # 确定基数表示
                        if is_unique:
                            cardinality = "1:1"
                        elif is_many_to_many:
                            cardinality = "N:M"
                        else:
                            cardinality = "1:N"
                        
                        relationship = Relationship(
                            from_table=table_name,
                            from_column=constrained_column,
                            to_table=referred_table,
                            to_column=referred_column,
                            relationship_type="one-to-many",
                            is_unique=is_unique,
                            is_many_to_many=is_many_to_many,
                            cardinality=cardinality,
                            description=f"{table_name}.{constrained_column} -> {referred_table}.{referred_column}"
                        )
                        relationships.append(relationship)
            
            return relationships
            
        except Exception as e:
            logger.error(f"获取表关系失败: {str(e)}")
            raise
    
    def close(self):
        """关闭数据库连接"""
        if self.engine:
            self.engine.dispose()
            logger.info("数据库连接已关闭")